package com.shujia.mllib

import org.apache.spark.ml.classification.LogisticRegressionModel
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo06ImagePredict {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo06ImagePredict")
      .getOrCreate()

    import spark.implicits._

    val imageDF: DataFrame = spark
      .read
      .format("image")
      .load("spark/data/mllib/data/images")

    // 如果需要使用模型对新的数据进行预测 则同样要经过相同的数据特征工程处理
    val predictImageDF: DataFrame = imageDF
      .select($"image.origin" as "path", $"image.data" as "data")
      .as[(String, Array[Byte])]
      .map(t2 => {
        val fileName: String = t2._1.split("/").reverse.head
        val newDataList: List[Double] = t2._2.toList.map(i => {
          if (i >= 0) {
            0.0
          } else {
            255.0
          }
        })
        (fileName, Vectors.dense(newDataList.toArray).toSparse)
      }).toDF("fileName", "features")

    val lRModel: LogisticRegressionModel = LogisticRegressionModel.load("spark/data/mllib/image")

    lRModel.transform(predictImageDF).show()


  }

}
